import numpy as np
import matplotlib.pyplot  as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.spatial.transform import Rotation

# 设置绘图参数
# 确保自己电脑装了 Latex, 用 Latex 渲染是为了好看
# 因为我的 Latex 没安装支持中文的包, 暂时懒得装, 就用英文了, 反正文字也不多, 凑合着用
plt.rcParams.update({
    'font.size': 14,
    'font.family': 'serif',
    'text.usetex': True,
    'text.latex.preamble': r'\usepackage{amsmath}'
})

# 莫兰迪色系, 我比较喜欢这个色系, 看着"温柔"一些
colors = {
    'red': (213/255, 126/255, 126/255),
    'green': (126/255, 178/255, 126/255),
    'blue': (126/255, 126/255, 178/255),
    'purple': (178/255, 126/255, 178/255),
    'gray': (178/255, 178/255, 178/255),
    'black': (50/255, 50/255, 50/255),
    'orange': (213/255, 178/255, 126/255),
    'yellow': (213/255, 213/255, 126/255),
    'cyan': (126/255, 213/255, 213/255),
    'brown': (150/255, 100/255, 70/255)
}

## 自带的 3D 箭头不好看, 所以不如自定义

def draw_3d_arrow(ax, start, end, color, arrow_scale=0.15, label=None):
    """绘制带3D圆锥箭头的向量"""
    start = np.array(start)
    end = np.array(end)
    direction = end - start
    length = np.linalg.norm(direction)
    if length == 0:
        return

    # 计算箭头长度和半径
    arrow_length = min(arrow_scale, length * 0.3)
    arrow_radius = arrow_length * 0.2

    # 绘制线段主体
    line_end = end - direction/length * arrow_length
    ax.plot([start[0], line_end[0]], [start[1], line_end[1]], [start[2], line_end[2]],
            color=color, linewidth=2)

    # 创建圆锥箭头
    # 生成圆锥底面圆
    theta = np.linspace(0, 2*np.pi, 20)
    r = arrow_radius
    h = arrow_length
    x = r * np.cos(theta)
    y = r * np.sin(theta)
    z = np.zeros_like(x) - h  # 圆锥底面在z=-h处

    # 旋转圆锥到箭头方向
    if not np.allclose(direction, [0, 0, 1]):
        rotation_axis = np.cross([0, 0, 1], direction)
        rotation_angle = np.arccos(np.dot([0, 0, 1], direction)/length)
        rotation = Rotation.from_rotvec(rotation_angle * rotation_axis/np.linalg.norm(rotation_axis))
        x, y, z = rotation.apply(np.vstack([x, y, z]).T).T

    # 平移圆锥到终点
    cone_x = x + end[0]
    cone_y = y + end[1]
    cone_z = z + end[2]

    # 圆锥顶点在终点位置
    apex = end

    # 创建圆锥侧面三角形
    triangles = []
    for i in range(len(theta)-1):
        triangles.append([apex, [cone_x[i], cone_y[i], cone_z[i]], [cone_x[i+1], cone_y[i+1], cone_z[i+1]]])

    # 绘制圆锥侧面
    ax.add_collection3d(Poly3DCollection(triangles, color=color, alpha=0.8))

    # 绘制圆锥底面
    ax.add_collection3d(Poly3DCollection([list(zip(cone_x, cone_y, cone_z))], color=color, alpha=0.5))

    # 添加标签
    if label:
        text_pos = end + direction * 0.1
        ax.text(text_pos[0], text_pos[1], text_pos[2], label, color=color)

def rotation_matrix_from_quaternion(q):
    """从四元数计算旋转矩阵"""
    q0, q1, q2, q3 = q
    return np.array([
        [q0**2 + q1**2 - q2**2 - q3**2, 2*(q1*q2 - q0*q3), 2*(q1*q3 + q0*q2)],
        [2*(q1*q2 + q0*q3), q0**2 - q1**2 + q2**2 - q3**2, 2*(q2*q3 - q0*q1)],
        [2*(q1*q3 - q0*q2), 2*(q2*q3 + q0*q1), q0**2 - q1**2 - q2**2 + q3**2]
    ])

def draw_rotation_arc_with_arrow(ax, center, axis, radius, angle, color, n_points=100):
    """绘制带3D圆锥箭头的旋转弧线, 显示接近360度的视觉效果"""
    # 创建垂直于旋转轴的平面
    if np.allclose(axis, [0, 0, 1]):
        ortho1 = np.array([1, 0, 0])
        ortho2 = np.array([0, 1, 0])
    else:
        ortho1 = np.random.randn(3)
        ortho1 -= ortho1.dot(axis) * axis
        ortho1 /= np.linalg.norm(ortho1)
        ortho2 = np.cross(axis, ortho1)

    # 调整弧线位置到靠近轴顶部
    arc_center = center + axis * 0.8  # 将弧线中心沿轴向上移动

    # 绘制接近360度的弧线 (视觉上是完整的圆环)
    theta_full = np.linspace(0, 2*np.pi, n_points)
    full_circle_points = []
    for t in theta_full:
        rot_vec = t * axis
        rotation = Rotation.from_rotvec(rot_vec)
        point = arc_center + radius * rotation.apply(ortho1)
        full_circle_points.append(point)
    full_circle_points = np.array(full_circle_points)

    # 绘制完整的圆形轨迹
    ax.plot(full_circle_points[:,0], full_circle_points[:,1], full_circle_points[:,2],
            color=color, linewidth=1, alpha=0.3, linestyle='--')

    # 绘制实际旋转角度的弧线
    theta = np.linspace(0, angle, n_points)
    points = []
    for t in theta:
        rot_vec = t * axis
        rotation = Rotation.from_rotvec(rot_vec)
        point = arc_center + radius * rotation.apply(ortho1)
        points.append(point)
    points = np.array(points)

    # 绘制实际的旋转弧线
    ax.plot(points[:,0], points[:,1], points[:,2],
            color=color, linewidth=2, alpha=0.8)

    # 添加3D圆锥箭头
    if len(points) >= 2:
        end_point = points[-1]
        prev_point = points[-2]
        direction = end_point - prev_point
        arrow_length = 0.2

        # 计算箭头终点位置
        arrow_end = end_point + direction/np.linalg.norm(direction) * arrow_length

        # 使用之前定义的3D箭头函数绘制箭头
        draw_3d_arrow(ax, end_point, arrow_end, color, arrow_scale=0.2)

    # 绘制角度标记 (在弧线外侧)
    mid_angle = angle/2
    rot_vec = mid_angle * axis
    rotation = Rotation.from_rotvec(rot_vec)
    text_pos = arc_center + (radius + 0.15) * rotation.apply(ortho1)

    ax.text(text_pos[0], text_pos[1], text_pos[2],
            f"{np.degrees(angle):.0f}°",
            color=color, fontsize=14, ha='center', va='center')

def create_cube(center, size=0.5):
    """创建3D立方体"""
    half = size / 2
    vertices = np.array([
        [center[0]-half, center[1]-half, center[2]-half], # 0
        [center[0]+half, center[1]-half, center[2]-half], # 1
        [center[0]+half, center[1]+half, center[2]-half], # 2
        [center[0]-half, center[1]+half, center[2]-half], # 3
        [center[0]-half, center[1]-half, center[2]+half], # 4
        [center[0]+half, center[1]-half, center[2]+half], # 5
        [center[0]+half, center[1]+half, center[2]+half], # 6
        [center[0]-half, center[1]+half, center[2]+half]  # 7
    ])

    # 定义立方体的面
    faces = [
        [vertices[0], vertices[1], vertices[2], vertices[3]], # 底面
        [vertices[4], vertices[5], vertices[6], vertices[7]], # 顶面
        [vertices[0], vertices[1], vertices[5], vertices[4]], # 前面
        [vertices[2], vertices[3], vertices[7], vertices[6]], # 后面
        [vertices[1], vertices[2], vertices[6], vertices[5]], # 右面
        [vertices[0], vertices[3], vertices[7], vertices[4]]  # 左面
    ]

    # 定义立方体的边 (12条边)
    edges = [
        [vertices[0], vertices[1]],
        [vertices[1], vertices[2]],
        [vertices[2], vertices[3]],
        [vertices[3], vertices[0]],
        [vertices[4], vertices[5]],
        [vertices[5], vertices[6]],
        [vertices[6], vertices[7]],
        [vertices[7], vertices[4]],
        [vertices[0], vertices[4]],
        [vertices[1], vertices[5]],
        [vertices[2], vertices[6]],
        [vertices[3], vertices[7]]
    ]

    return vertices, faces, edges

def plot_3d_rotation():
    """绘制3D旋转示意图"""
    fig = plt.figure(figsize=(16, 14))

    # 创建网格布局
    gs = fig.add_gridspec(2, 1, height_ratios=[3, 1])
    ax = fig.add_subplot(gs[0], projection='3d')
    text_ax = fig.add_subplot(gs[1])
    text_ax.axis('off')

    # 定义初始坐标系和点
    original_points = np.array([
        [0, 0, 0],
        [1, 0, 0],
        [0, 1, 0],
        [0, 0, 1]
    ])

    # 定义四元数表示的旋转 (绕轴[1,1,1]旋转60度)
    theta = np.pi/3   # 60度
    axis = np.array([1, 1, 1])
    axis = axis / np.linalg.norm(axis)  # 归一化

    # 计算对应的四元数
    q = np.array([
        np.cos(theta/2),           # q0
        axis[0]*np.sin(theta/2),   # q1
        axis[1]*np.sin(theta/2),   # q2
        axis[2]*np.sin(theta/2)    # q3
    ])

    # 从四元数计算旋转矩阵
    R = rotation_matrix_from_quaternion(q)

    # 应用旋转
    rotated_points = np.dot(original_points, R.T)

    # 创建并旋转立方体
    cube_vertices, cube_faces, cube_edges = create_cube([0.7, 0.7, 0.7], size=0.8)
    rotated_cube_vertices = np.dot(cube_vertices, R.T)
    rotated_cube_faces = []
    for face in cube_faces:
        rotated_face = np.dot(np.array(face), R.T)
        rotated_cube_faces.append(rotated_face)

    # 旋转立方体的边
    rotated_cube_edges = []
    for edge in cube_edges:
        rotated_edge = np.dot(np.array(edge), R.T)
        rotated_cube_edges.append(rotated_edge)

    # 绘制原始坐标系
    draw_3d_arrow(ax, [0, 0, 0], [1, 0, 0], colors['red'], 0.1)
    draw_3d_arrow(ax, [0, 0, 0], [0, 1, 0], colors['green'], 0.1)
    draw_3d_arrow(ax, [0, 0, 0], [0, 0, 1], colors['blue'], 0.1)
    ax.text(1.1, 0, 0, 'X', color=colors['red'], fontsize=14)
    ax.text(0, 1.1, 0, 'Y', color=colors['green'], fontsize=14)
    ax.text(0, 0, 1.1, 'Z', color=colors['blue'], fontsize=14)

    # 绘制旋转后的坐标系
    rotated_axes = np.dot(np.eye(3), R.T)
    draw_3d_arrow(ax, [0, 0, 0], rotated_axes[0], colors['red'], 0.1)
    draw_3d_arrow(ax, [0, 0, 0], rotated_axes[1], colors['green'], 0.1)
    draw_3d_arrow(ax, [0, 0, 0], rotated_axes[2], colors['blue'], 0.1)
    ax.text(rotated_axes[0,0], rotated_axes[0,1], rotated_axes[0,2], "X'", color=colors['red'], fontsize=14)
    ax.text(rotated_axes[1,0], rotated_axes[1,1], rotated_axes[1,2], "Y'", color=colors['green'], fontsize=14)
    ax.text(rotated_axes[2,0], rotated_axes[2,1], rotated_axes[2,2], "Z'", color=colors['blue'], fontsize=14)

    # 绘制旋转轴

    draw_3d_arrow(ax, [0, 0, 0], axis + 0.6, colors['purple'], 0.1, label=r'$\text{Rotation Axis }\pmb{n}$')
    ax.text(axis[0]*2, axis[1]*2, axis[2]*1.8, r'$\pmb{n}=(n_x,n_y,n_z)$', color=colors['purple'], fontsize=14)

    # 绘制立方体及其旋转
    # 绘制原始立方体 (半透明)
    ax.add_collection3d(Poly3DCollection(cube_faces, facecolors=colors['green'], linewidths=0, edgecolors='none', alpha=0.2))

    # 绘制旋转后的立方体 (半透明)
    ax.add_collection3d(Poly3DCollection(rotated_cube_faces, facecolors=colors['yellow'], linewidths=0, edgecolors='none', alpha=0.3))

    # 绘制立方体的边
    for edge in cube_edges:
        ax.plot([edge[0][0], edge[1][0]], [edge[0][1], edge[1][1]], [edge[0][2], edge[1][2]],
                color=colors['brown'], linestyle='--', linewidth=1.5)

    for edge in rotated_cube_edges:
        ax.plot([edge[0][0], edge[1][0]], [edge[0][1], edge[1][1]], [edge[0][2], edge[1][2]],
                color=colors['brown'], linestyle='-', linewidth=1.5)

    # 绘制旋转弧线(带箭头)
    cube_center = np.mean(cube_vertices, axis=0)
    rotated_cube_center = np.mean(rotated_cube_vertices, axis=0)
    draw_rotation_arc_with_arrow(ax, cube_center, axis, 0.4, theta, colors['orange'])

    # 标记四元数分量
    quaternion_text = r'$\pmb{q}=\begin{bmatrix}q_0\\q_1\\q_2\\q_3\end{bmatrix}=\begin{bmatrix}' + \
                      f'{q[0]:.2f}\\\\{q[1]:.2f}\\\\{q[2]:.2f}\\\\{q[3]:.2f}' + r'\end{bmatrix}$'
    ax.text(-0.5, -0.5, 1.8, quaternion_text, color=colors['black'], fontsize=14)

    # 设置坐标轴范围
    ax.set_xlim([-0.5, 1.5])
    ax.set_ylim([-0.5, 1.5])
    ax.set_zlim([-0.5, 1.5])
    ax.set_xlabel('X', fontsize=14)
    ax.set_ylabel('Y', fontsize=14)
    ax.set_zlabel('Z', fontsize=14)
    ax.set_title('3D  Rotation Visualization with Quaternion', pad=20, fontsize=16)

    # 添加公式说明
    formula_text1 = r"$\text{1. Rigid transformation: }\pmb{a}=\pmb{T}\pmb{b},\pmb{T}\in\mathrm{SE}(3)$"
    formula_text2 = r"2. Rotation matrix from quaternion:"
    formula_text3 = r"$\pmb{R}=\pmb{v}\pmb{v}^{\mathrm{T}} + s^{2}\pmb{I} + 2s\pmb{v}^{\wedge} + (\pmb{v}^{\wedge})^{2}=" + \
                    r"\begin{bmatrix} q_0^2 + q_1^2 - q_2^2 - q_3^2 & 2(q_1 q_2 - q_0 q_3) & 2(q_1 q_3 + q_0 q_2) \\" + \
                    r"2(q_1 q_2 + q 0 q_3) & q_0^2 - q_1^2 + q_2^2 - q_3^2 & 2(q_2 q_3 - q_0 q_1) \\" + \
                    r"2(q_1 q_3 - q_0 q_2) & 2(q_2 q_3 + q_0 q_1) & q_0^2 - q_1^2 - q_2^2 + q_3^2 \end{bmatrix}$"
    formula_text4 = r"$\text{3. Quaternion representation: }\pmb{q}=\Vert\pmb{q}\Vert\left(\cos{\frac{\theta}{2}}+\pmb{n}\sin{\frac{\theta}{2}}\right)$"
    formula_text5 = r"$\text{where }\theta \text{ is the rotation angle, and } \pmb{n} \text{ is the rotation axis (unit vector)}$"

    text_ax.text(0.02, 0.2, formula_text1 + "\n\n" + formula_text2 + "\n" + formula_text3 + "\n\n" + formula_text4 + "\n" + formula_text5,
                fontsize=14, bbox=dict(facecolor='white', alpha=0.8))

    plt.tight_layout()
    plt.show()

plot_3d_rotation()
